import "babel-polyfill";
import {enumerate, map, filter, range} from "../../../algorithm/iterator";
import Integer from "../../number/integer";
import {invmod} from "../../number-theory/gcd";
import PolynomialZn from "../polynomial-Zn";
import {gcd} from "../gcd";

var squareFree = function(polynomial) {
    var PolynomialType = polynomial.constructor;
    var instance = PolynomialType.instance;
    var prime = PolynomialType.modulus;
    var result = [[PolynomialType.one, 0]];

    var f = polynomial;
    var h = PolynomialType.zero;
    for (let degree = 0; ; ++degree) {
        let d = f.deriv();
        let f1 = PolynomialType::gcd(f, d);
        let h1 = instance.divide(f, f1);
        let g = instance.divide(h, h1);
        if (degree && g.degree >= 0) {
            result.push([g, degree]);
        }
        f = f1;
        h = h1;
        if (h.degree === 0) {
            break;
        }
    }

    if (f.degree !== 0) {
        prime = prime.lastLimb;
        let s_poly = new PolynomialType(...range(0, f.degree + 1, prime)::map(degree => f.get(degree)));
        let k = result.length;
        for (let [j, [p, d]] of squareFree(s_poly)::enumerate()) {
            result.push([p, d * prime]);
            let head = (j + 1) * k;
            for (let i = 1; i < k; ++i) {
                let g = PolynomialType::gcd(result[i][0], result[head][0]);
                result.push([g, result[i][1] + result[head][1]]);
                result[i][0] = instance.divide(result[i][0], g);
                result[head][0] = instance.divide(result[head][0], g);
            }
        }
    }

    return result::filter(([poly]) => poly.degree > 0);
};

var fast_pow = function(polynomial, modulus) {
    var PolynomialType = polynomial.constructor;
    var degree = polynomial.degree;
    if (degree === -Infinity) {
        return PolynomialType.zero;
    }
    var prime = PolynomialType.modulus.lastLimb;
    var result = [polynomial.get(0)];
    for (let d = 1; d <= degree; ++d) {
        for (let i = 1; i < prime; ++i) {
            result.push(Integer.zero);
        }
        result.push(polynomial.get(d));
    }
    return PolynomialType.instance.mod(new PolynomialType(...result), modulus);
};

var distinctDegree = function*(polynomial) {
    var PolynomialType = polynomial.constructor;
    var instance = PolynomialType.instance;
    var prime = PolynomialType.modulus;
    var f = polynomial;
    var base = new PolynomialType(0, 1);
    var h = base;
    var pow = Integer.compare(prime, Integer[3]) <= 0
        ? p => fast_pow(p, polynomial)
        : p => instance.pow(p, prime, polynomial);
    for (let degree = 1; ; ++degree) {
        if (f.degree === 0) {
            return;
        } else if (degree << 1 > f.degree) {
            yield [f, f.degree];
            return;
        } else {
            h = pow(h);
            let g = PolynomialType::gcd(instance.subtract(h, base), f);
            if (g.degree > 0) {
                yield [g, degree];
                f = instance.divide(f, g);
            }
        }
    }
};

var equalDegree = function*(polynomial, degree) {
    var PolynomialType = polynomial.constructor;
    var instance = PolynomialType.instance;
    var prime = PolynomialType.modulus;
    var stack = [polynomial];
    if (Integer.eq(prime, Integer[2])) {
        var pow = function(p, modulus) {
            var q = p;
            for (let i = 1; i < degree; ++i) {
                p = fast_pow(p, modulus);
                q = instance.add(q, p);
            }
            return q;
        };
    } else {
        let exp = Integer.shr(Integer.pow(prime, degree), 1);
        var pow = (p, modulus) => instance.pow(p, exp, modulus);
    }

    while (stack.length) {
        let top = stack.pop();
        let top_degree = top.degree;
        if (top_degree === degree) {
            yield top;
            continue;
        }
        while (true) {
            let a = new PolynomialType();
            while (a.degree === -Infinity) {
                a = new PolynomialType(...range(top_degree)::map(x => Integer.random(prime)));
            }
            let g = PolynomialType::gcd(a, top);
            if (g.degree > 0) {
                stack.push(g, instance.divide(top, g));
                break;
            }
            let b = instance.subtract(pow(a, top), PolynomialType.one);
            g = PolynomialType::gcd(b, top);
            if (g.degree > 0 && g.degree < top_degree) {
                stack.push(g, instance.divide(top, g));
                break;
            }
        }
    }
};

export var factorUnique = function*(polynomial) {
    for (let [p, degree] of distinctDegree(this.instance.from(polynomial))) {
        yield* equalDegree(p, degree);
    }
};

export default function* factor(polynomial) {
    polynomial = this.instance.from(polynomial);
    if (polynomial.degree === -Infinity) {
        yield [this.zero, 1];
        return;
    }
    var leading_coeff = polynomial.leadingCoeff;
    if (!Integer.eq(leading_coeff, Integer.one)) {
        yield [new this(leading_coeff), 1];
    }
    for (let [p, degree] of squareFree(polynomial)) {
        yield* this::factorUnique(p)::map(f => [f, degree]);
    }
};
